
#ifndef _MLS_KERNEL_CUDA_H_
#define _MLS_KERNEL_CUDA_H_

#define MIN(a,b) ((a)>(b)?(b):(a))
#define MAX(a,b) ((a)<(b)?(b):(a))
#define SIGN(a)   (a<0)?(-1):(1)
#define dot(A,B) (A[0]*B[0] + A[1]*B[1] + A[2]*B[2])
#define cross(A,B) make_float3(A.y*B.z-B.y*A.z, B.x*A.z-A.x*B.z, A.x*B.y-B.x*A.y)

/// Constant values for cubic equation computation
const float K1DIV3   = 0.3333333;    /// 1/3
const float K1DIV9   = 0.1111111;    /// 1/9
const float K1DIV54  = 0.0185185;    /// 1/54
const float EPSILON  = 1.0E-5;       

texture<float4, 1, cudaReadModeElementType> texQ; /// control points texture
texture<float4, 1, cudaReadModeElementType> texP; /// control points texture
texture<float4, 2, cudaReadModeElementType> texPositions; /// positions texture
texture<float4, 2, cudaReadModeElementType> texNormals; /// positions texture

/// Returns a root of the cubic equation x^3 + ax^2 + bx + c = 0
/// @param a coefficient of the second degree term
/// @param b coefficient of the first degree term
/// @param c coefficient of the independent term
/// @return a one root of the cubic equation
__device__ float cubicroot(const float a, const float b, const float c) {
   
   const float Q = (3.0*b - a*a)*K1DIV9;
   const float R = (9.0*a*b - 27.0*c - 2.0*a*a*a)*K1DIV54;
   const float D = Q*Q*Q + R*R;
   const float sqrtD = sqrt(D);
   
   if (D>=0.0) {  /// one real, two complex :: D==0 root reals, at least two are equal 
      const float S = SIGN(R+sqrtD)*pow(abs(R+sqrtD), K1DIV3);
      const float T = SIGN(R-sqrtD)*pow(abs(R-sqrtD), K1DIV3);
      return -a*K1DIV3 + S + T;
   }
   else {   /// all root are real and unequal
      const float theta = acos(R/sqrt(-Q*Q*Q));
      return 2.0*sqrt(-Q)*cos(theta*K1DIV3)-a*K1DIV3;
   }
}

/// Returns the largest root of a depressed quartic equation  
/// defined by the 3x3 MLS-based correlation matrix M
/// @param M0 first row of matrix M
/// @param M1 second row of matrix M
/// @param M2 third row of matrix M
/// @return the largest root
__device__ float largestRoot(const float *M0, const float *M1, const float *M2) {
   
   const float detM = M0[0]*(M1[1]*M2[2]-M1[2]*M2[1])+
                      M0[1]*(M1[2]*M2[0]-M1[0]*M2[2])+
                      M0[2]*(M1[0]*M2[1]-M2[0]*M1[1]);
   
   /// setting the cubic equation
   float normaM[3], normaM2=0.0, prodNorma2=0.0;
   normaM[0] = dot(M0,M0);
   normaM2 += normaM[0];
   normaM[1] = dot(M1,M1);
   normaM2 += normaM[1];
   normaM[2] = dot(M2,M2);
   normaM2 += normaM[2];
   
   prodNorma2 += normaM[0]*normaM[1];
   prodNorma2 += normaM[1]*normaM[2];
   prodNorma2 += normaM[2]*normaM[0];
   
   float prodesc, prodesc2=0.0;
   prodesc = dot(M0,M1);
   prodesc2 += prodesc*prodesc;
   prodesc = dot(M1,M2);
   prodesc2 += prodesc*prodesc;
   prodesc = dot(M2,M0);
   prodesc2 += prodesc*prodesc;
   
   const float a = -4.0*normaM2;
   const float b = 16.0*(prodNorma2-prodesc2);
   const float c = -64.0*detM*detM;
   /// getting the cubic root
   float p2 = cubicroot(a,b,c);
   const float _p = sqrt(p2);
   const float minusp2_plus_4M2 = -p2 + 4.0*normaM2;
   const float _16detMdivp = (16.0*detM)/_p;
   
   /// the largest quartic root
   float maxroot = MAX(-_p + sqrt(minusp2_plus_4M2 - _16detMdivp), 
                        _p + sqrt(minusp2_plus_4M2 + _16detMdivp));
   
   return maxroot*0.5;
}

/// Computes the rotation vector
/// @param M0 first row of correlation matrix
/// @param M1 second row of correlation matrix
/// @param M2 third row of correlation matrix
/// @param eigenvalue 
/// @return a vector that defines the optimal rotation
__device__ float3 rotationVector(const float *M0, const float *M1, const float *M2, const float eigenvalue) {
   
   /// solving the 3x3 linear system: (N-eigenvalue*I) e = V
   // the right side
   float V[3]; 
   V[0] = M1[2]-M2[1]; V[1] = M2[0]-M0[2]; V[2] = M0[1]-M1[0];
   // the left side
   float N0[3], N1[3], N2[3];
   N0[0]=M0[0]+M0[0]; N0[1]=M0[1]+M1[0]; N0[2]=M0[2]+M2[0];
   N1[0]=M1[0]+M0[1]; N1[1]=M1[1]+M1[1]; N1[2]=M1[2]+M2[1];
   N2[0]=M2[0]+M0[2]; N2[1]=M2[1]+M1[2]; N2[2]=M2[2]+M2[2];
   N0[0] -= eigenvalue;
   N1[1] -= eigenvalue;
   N2[2] -= eigenvalue;
   
   /// Solve the 3x3 system
   const float sigma = (N0[0]*N2[1] - N2[0]*N0[1])/(N0[0]*N1[1] - N1[0]*N0[1]);
   const float vez = (N0[0]*V[2]  -N2[0]*V[0]  -sigma*(N0[0]*V[1] - N1[0]*V[0]))/
                     (N0[0]*N2[2] -N2[0]*N0[2]-sigma*(N0[0]*N1[2]-N1[0]*N0[2]));
   const float vey = ((N0[0]*V[1] -N1[0]*V[0])-(N0[0]*N1[2]-N1[0]*N0[2])*vez)/
                     (N0[0]*N1[1] -N1[0]*N0[1]);
   const float vex = (V[0]-N0[2]*vez-N0[1]*vey)/N0[0];
   
   return make_float3(vex, vey, vez);
}

/// Returns the weight associated with the i-th control point
/// @param p the point being evaluated
/// @param Pi the i-th control point
/// @return the weight value
__device__ float weight(const float4 &p, const float4 &Pi) {
   
   float vd[3];
   vd[0] = Pi.x-p.x; vd[1] = Pi.y-p.y; vd[2] = Pi.z-p.z;
   const float d = sqrt(dot(vd,vd));
   if (d==0) return 1e8;
   else return 1.0/(d*d);
}

/// Computes the weighted centroids of control points
/// @param p the point being evaluated
/// @return pstar weighted centroid defined by the initial control points
/// @return qstar weighted centroid defined by the deformed control points
/// @param P set of initial control points
/// @param Q set of deformed control points
/// @param NCTRLPOINTS number of control points
__device__ void weightedCentroids(const float4 &p, float *pstar, float *qstar, const int NCTRLPOINTS) {
   
   float Wt=0.0, W;
   pstar[0]=pstar[1]=pstar[2]=0.0;
   qstar[0]=qstar[1]=qstar[2]=0.0;
   for (uint i=0; i<NCTRLPOINTS; i++) {
      float4 P = tex1Dfetch(texP, i);
      float4 Q = tex1Dfetch(texQ, i);
      W = weight(p,P);
      Wt += W;
      
      pstar[0] += W*P.x;
      pstar[1] += W*P.y;
      pstar[2] += W*P.z;
      
      qstar[0] += W*Q.x;
      qstar[1] += W*Q.y;
      qstar[2] += W*Q.z;
   }
   pstar[0] = pstar[0]/Wt;   pstar[1] = pstar[1]/Wt;   pstar[2] = pstar[2]/Wt;
   qstar[0] = qstar[0]/Wt;   qstar[1] = qstar[1]/Wt;   qstar[2] = qstar[2]/Wt;
}

/// Builds the correlation matrix M
/// @param p the point being evaluated
/// @param pstar initial weighted centroid
/// @param qstar deformed weighted centroid
/// @param P set of initial control points
/// @param Q set of deformed control points
/// @param NCTRLPOINTS number of control points
/// @return M0 first row of correlation matrix
/// @return M1 second row of correlation matrix
/// @return M2 third row of correlation matrix
__device__ void correlationMatrix(const float4 &p, const float *pstar, const float *qstar,
                                  const int NCTRLPOINTS, float *M0, float *M1, float *M2) {
   
   float3 pjhat, qjhat;
   M0[0]=M0[1]=M0[2]=0.0;
   M1[0]=M1[1]=M1[2]=0.0;
   M2[0]=M2[1]=M2[2]=0.0;
   for (int j=0; j<NCTRLPOINTS; j++) {
      float4 P = tex1Dfetch(texP, j);
      float4 Q = tex1Dfetch(texQ, j);
      float W = weight(p,P);
      pjhat.x = P.x-pstar[0];  pjhat.y = P.y-pstar[1]; pjhat.z = P.z-pstar[2];
      qjhat.x = Q.x-qstar[0];  qjhat.y = Q.y-qstar[1]; qjhat.z = Q.z-qstar[2];
      
      M0[0] += W*qjhat.x*pjhat.x; M0[1] += W*qjhat.x*pjhat.y; M0[2] += W*qjhat.x*pjhat.z;
      M1[0] += W*qjhat.y*pjhat.x; M1[1] += W*qjhat.y*pjhat.y; M1[2] += W*qjhat.y*pjhat.z;
      M2[0] += W*qjhat.z*pjhat.x; M2[1] += W*qjhat.z*pjhat.y; M2[2] += W*qjhat.z*pjhat.z;
   }
}

/// The MLS kernel. 
/// Given the control points sets this kernel computes an optimal rigid-body 
/// transformation to find the deformed the position of a 3D point 
/// @param pos 
/// @param pos 
__global__ void MLSkernel(float4* pos, float4* normals, unsigned int width, 
unsigned int npositions, const int NCTRLPOINTS) {
   
   /// calculate vertex coordinates
   unsigned int x = blockIdx.x*blockDim.x + threadIdx.x;
   unsigned int y = blockIdx.y*blockDim.y + threadIdx.y;
   float4 p = tex2D(texPositions, x, y);
   float4 n = tex2D(texNormals, x, y);
   
   if (y*width+x > npositions) {
      pos[y*width+x] = p;
      normals[y*width+x] = n;
   } else {
   /// compute the MLS deformation
      /// compute weighted centroids
      float pstar[3], qstar[3];
      weightedCentroids(p, pstar, qstar, NCTRLPOINTS);
      
      /// compute the correlation matrix
      float M0[3], M1[3], M2[3]; /// the matrix rows
      correlationMatrix(p, pstar, qstar, NCTRLPOINTS, M0, M1, M2);
      
      /// compute the largest root of polynomial and the eigenvalue
      const float traceM = M0[0] + M1[1] + M2[2];
      const float eigenvalue = largestRoot(M0,M1,M2) + traceM;
      
      /// compute the rotation vector
      float3 u = rotationVector(M0,M1,M2, eigenvalue);
      
      /// modify the vertex position (apply rotation defined by u)
      float3 vhat = make_float3(p.x-pstar[0], p.y-pstar[1], p.z-pstar[2]);
      
      /// rotate the vertex position
      const float3 tmpp1 = cross(u,vhat);
      const float3 tmpp2 = cross(tmpp1, u);
      const float numpx = 2.0*(tmpp2.x - tmpp1.x);
      const float numpy = 2.0*(tmpp2.y - tmpp1.y);
      const float numpz = 2.0*(tmpp2.z - tmpp1.z);
      
      const float denom = 1.0 + (u.x*u.x + u.y*u.y + u.z*u.z);
      const float vdx = vhat.x - numpx/denom;
      const float vdy = vhat.y - numpy/denom;
      const float vdz = vhat.z - numpz/denom;
      
      const float vrx = qstar[0] + vdx;
      const float vry = qstar[1] + vdy;
      const float vrz = qstar[2] + vdz;
      
      pos[y*width+x] = make_float4(vrx,vry,vrz,1.0);
      
      /// rotate the normal vector
      const float3 tmpn1 = cross(u,n);
      const float3 tmpn2 = cross(tmpn1,u);
      const float numnx = 2.0*(tmpn2.x - tmpn1.x);
      const float numny = 2.0*(tmpn2.y - tmpn1.y);
      const float numnz = 2.0*(tmpn2.z - tmpn1.z);
      
      const float nrx = n.x - numnx/denom;
      const float nry = n.y - numny/denom;
      const float nrz = n.z - numnz/denom;
      
      normals[y*width+x] = make_float4(nrx,nry,nrz,1.0);
   }
}

#endif
